"""Isolated shock problem"""

from hogs.grids import grid1d as grid
from hogs.solvers.euler1d import EESolver
from hogs.solvers.flux.euler.eulerian_flux import GodunovFluxEulerian, \
     EulerLxf, WAFFluxEulerian

from hogs.solvers.time_step_functions import EulerEquationsTimeStep
from hogs.solvers.primitive_variable_functions import EEPrimitiveVariable

import numpy

gamma = 1.4

# create and initialize the grid
g = grid.Grid1D()
g.initialize(xlow=-1.0, xhigh=1.0, dx=0.002, nb=2, nvar=3)

# construct the solver
solver = EESolver(gamma=1.4, tf=0.15, grid=g, nvar=3 )

# set the flux function
#solver.flux_function = flux.WAFFluxEulerian(gamma=1.4)
solver.flux_function = EulerLxf()

# set the grid for the flux function
solver.flux_function.set_grid( solver.grid )

# set the time step function
solver.time_step_function = EulerEquationsTimeStep(grid=solver.grid)

# primitive variable function
solver.primitive_variable_function = EEPrimitiveVariable(gamma=gamma,
                                                         grid=solver.grid)

# process command line 
solver.setup()

# set the variables
grid = solver.grid

# cell centers
x = grid.xc
q = grid.q

gamma = 1.4
gm1bgp1 = (gamma-1.0)/(gamma + 1.0)
gp1b2g = 0.5 * (gamma + 1.0)/gamma
gm1b2g = 0.5 * (gamma - 1.0)/gamma

# preshock values and shock speed
S = 7.0/3.0
rhor = 1.0; ur = 0.0; pr = 1.0

# pressure ratio (pl/pr)
ar = numpy.sqrt(gamma * pr/rhor)
pressure_ratio = (( (S - ur)/ar )**2 - gm1b2g)/gp1b2g

pl = pressure_ratio * pr

# density ratio (rhol/rhor)
density_ratio = (pressure_ratio + gm1bgp1)/(pressure_ratio*gm1bgp1 + 1.0)
rhol = density_ratio * rhor

# velocity ratio
ul = S * (1.0 - 1.0/density_ratio) + ur/density_ratio

# Internal energy ratio
el = pl/(0.4 * rhol)
er = pr/(0.4 * rhor)

for i, j in enumerate(x):
    if j < 0.0:

        q[0,i] = rhol
        q[1,i] = rhol*ul

        q[2,i] = rhol * (0.5 * ul**2 + el)

    else:

        q[0,i] = rhor
        q[1,i] = rhor*ur

        q[2,i] = rhor * (0.5 * ur**2 + er)

solver.solve()
